package yyds.sniarbtej;

import yyds.sniarbtej.transformer.ClassTransformer;
import yyds.sniarbtej.utils.DebugUtil;

import java.lang.instrument.ClassFileTransformer;
import java.security.ProtectionDomain;
import java.util.ArrayList;
import java.util.List;
import java.util.Set;
import java.util.TreeSet;

public final class MyTransformer implements ClassFileTransformer {
    private final Environment environment;
    private final List<ClassTransformer> transformers = new ArrayList<>();
    final Set<String> transformedClass = new TreeSet<>();

    public MyTransformer(Environment environment) {
        this.environment = environment;
    }

    public void addTransformer(ClassTransformer transformer) {
        if (transformer != null) {
            if (this.environment.isAttachMode()) {
                System.out.println("Transformer: " + transformer.getClass().getName() + " is set to not load in attach mode, ignored.");
            } else {
                synchronized (this) {
                    this.transformers.add(transformer);
                    this.transformedClass.add(transformer.canTransformClassName().replace('/', '.'));
                }
            }
        }
    }

    public Set<String> getTransformedClass() {
        return this.transformedClass;
    }

    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined,
                            ProtectionDomain protectionDomain, byte[] classfileBuffer) {
        if (className == null) {
            return classfileBuffer;
        } else {
            for (ClassTransformer transformer : this.transformers) {
                if (transformer.canTransform(className)) {
                    try {
                        classfileBuffer = transformer.rewrite(classfileBuffer);
                    } catch (Exception e) {
                        DebugUtil.log("ERROR: " + e.getMessage());
                        e.printStackTrace();
                    }
                }
            }

            return classfileBuffer;
        }
    }
}
